import hilo, parser_generator, grammar_parser, sys, time
from distorm3 import Decode, Decode32Bits
from di3_interp import interpret

# create a parser for output from distorm3
di3_out_grammar_file = './di3_out.grm'
di3parser = parser_generator.makeParser(grammar_parser.parse(open(di3_out_grammar_file).read()))

callCounter = 1 # can only trace one function at a time (scoping issues with Python 2.7)

# 
# Represents an open and attached process.
# Provides some operations:
#   memory derefencing (derefence)
#   making anonymous functions for remote calls (remoteFunction)
# Requires (kind of) that detach() is called before exiting
# Static methods:
#   getRunningProcesses( )
#      Returns a list of running processes
#
class Process:
    def __init__( self, name = None, pid = 0 ):
        self.pid = pid # process id
        self.hProc = None # handle to the (open)
        self.open = False
        if ( name != None or pid != 0 ):
            self.attach( name )
    def derefInt( self, address, size ):
        return DerefInt( self, address, size )
    def attach( self, name ):
        if ( self.open ):
            raise Exception()
        if name != None:
            self.pid = hilo.getPid( name )
        self.hProc = hilo.openProcess( self.pid )
        self.open = True
    def detach( self ):
        if ( self.open ):
            hilo.closeHandle( self.hProc )
            self.pid = 0
            self.hProc = None
            self.open = False
    def writeInt( self, address, length, val ):
        if ( self.open ):
            hilo.writeInt( self.hProc, address, length, val )
    def writeMemory( self, address, byteArray ):
        if ( self.open ):
            hilo.writeMemory( self.hProc, address, len(byteArray), byteArray )
    def readMemory( self, address, length ):
        if ( self.open ):
            return hilo.readMemory( self.hProc, address, length )
    def readInt( self, address, length ):
        if ( self.open ):
            return hilo.readInt( self.hProc, address, length )
    def isAttached( self ):
        return self.open
    def genCallbackTrace( self, code, grabAmount = 10 ):       
        def shouldContinueDebugging( curInstr ):
            global callCounter
            if curInstr[2][0:3].lower() == "ret":
                callCounter -= 1
            elif curInstr[2][0:4].lower() == "call":
                callCounter += 1
            return callCounter
        
        # code is another function that takes an instruction and context and instruction counter
        # callbacks take: instr. address, context, and an instruction counter
        def fn( address, context, instrCounter ):
            # grab the code
            instructions = []
            hexString = ""
            grabbedAmount = grabAmount
            # guess the instruction length
            while len(instructions) <= 1:
                hexString = str( self.readMemory( address, grabbedAmount ) ).encode('hex')
                instructions = Decode( address, hexString.decode('hex'), Decode32Bits )
                grabbedAmount += 5
            curInstr = instructions[0] # use current instruction

            # determine whether to continue or not
            continueVal = shouldContinueDebugging( curInstr )

            if continueVal != 0:
                code( curInstr, context, instrCounter ) # call our code if we continue
            return continueVal
        return fn
    def remoteFunction( self, address, argFormat ):
        if ( self.open == False):
            raise Exception()
        def fn( args, tf = 0):
            hThread = hilo.makeRemoteThread( self.hProc, tf )
            cContext = hilo.getThreadContext( hThread )
            # place exitthread arg onto stack (to avoid segfaulting)
            cContext.Esp -= 4
            self.writeInt( cContext.Esp, 4, 0 )

            codeCave = hilo.virtualAlloc( self.hProc, 32 )
            myCodecave = codeCave
            self.writeInt( codeCave, 1, 0xE9 )
            self.writeInt( codeCave + 1, 4, address - 5 - codeCave )

            # set next instruction to execute
            cContext.Eip = codeCave
            
            # modify context based on argObject
            for reg in argFormat.regActions:
                setattr( cContext, reg[0].upper() + reg[1:], args[argFormat.regActions[reg]] )

            for (length,name) in argFormat.stackActions:
                cContext.Esp -= length
                self.writeInt( cContext.Esp, length, args[name] )
            # push a return address onto the stack (ExitThread)
            cContext.Esp -= 4
            self.writeInt( cContext.Esp, 4, hilo.getDefaultReturnAddress() )

            hilo.setThreadContext( hThread, cContext )
            hilo.resumeAndWait( hThread )
            hilo.virtualFree( self.hProc, 32, codeCave )
            hilo.closeHandle( hThread )
            
        return fn
    def determineNextInstruction(self, instr, context):
        continueVal = (context.Eip != hilo.getDefaultReturnAddress())
        return continueVal
    def traceIntoCall( self, remoteFunction, argFormat, argsToFunc, outputFileName, address = 0 ):
        # note, when address is not 0, remoteFunction simply waits until tracing is done, else it should be a call to a remote function
        # this behaviour is taken care of within this function
        global callCounter
        callCounter = 1
        traceStarted = False
        outputFile = open(outputFileName, "w")

        outputFile.write('get(I,A,V) :- put(I,A,V);C is I-1,get(C,A,V).\n')
        outputFile.write('put(0,esp,0).\n')
        outputFile.write('put(-1,X,undef(X)).\n\n')

        for reg in argFormat.regActions:
            regVal = argFormat.regActions[reg]
            outputFile.write('put(0,%s,%s).\n' % (reg,str(regVal)))            

        esp = 4
        argFormat.stackActions.reverse()        
        for (length,name) in argFormat.stackActions:
            outputFile.write('put(0,deref(%d),%s).\n' % (esp,name))
            esp += length
        argFormat.stackActions.reverse()

        def reformatInstr( instrObj ):
            strInstr = instrObj[2].lower()
            startInd = strInstr.find(',')
            if startInd == -1:
                startInd += len(strInstr)
            while strInstr[startInd] != ' ' and startInd >= 0:
                startInd -= 1
            if startInd > 0:
                save = strInstr[startInd:]
                rep = strInstr[0:startInd].replace(' ', '')
                strInstr = rep + save

            return strInstr
        
        def analysisFun( instr, context, instrCounter ):
            #traceStarted = True
            #print instr[2]
            refStr = reformatInstr(instr)
            #print refStr
            ast = di3parser.parse(refStr)
            #print str(ast)
            #outputFile.write( str(instrCounter) + "\t " + hex(instr[0]) + "\t " + str(ast) + "\n" )
            outputFile.write( str(interpret( ast, instrCounter )) + "\n" )      
            
        hilo.traceIntoCall( self.hProc, self.genCallbackTrace( analysisFun ), remoteFunction, argsToFunc, address )
        outputFile.close( )
    @staticmethod
    def getRunningProcesses( ):
        listProcesses = hilo.getProcesses( ) # list of tuples (image,pid)
        retDict = {}
        for (imageName,pid) in listProcesses: # build dictionary of lists
            if imageName in retDict:
                retDict[imageName] += [pid]
            else:
                retDict[imageName] = [pid]
        return retDict

#
# Defines a way to move argument values in Python to a context for a thread in C++
# Returns calling object to enable call chaining.
#
class ArgFormat:
    def __init__( self ):
    # format for stackActions elements:
    #   tuples: (length,name)
    # * regActions:
    #   associate regName -> argName
        self.stackActions = []
        self.regActions = {}
    def setReg( self, regName, argName ):
        self.regActions[regName] = argName
        return self
    def pushVal( self, length, argName ):
        self.stackActions.append( (length, argName) )
        return self
    def push( self, length, argName ):
        self.stackActions.append( (length, argName) )
        return self

class DerefInt:
    def __init__( self, process, address, size ):
        self.process = process
        self.address = address
        self.size = size
    def write( self, val ):
        self.process.writeInt( self.address, self.size, val )
    def read( self ):
        return (self.process).readInt( self.address, self.size )

